#!/usr/bin/env python3
"""
通用关键字段信息抽取 - 部署脚本

使用 vLLM 和 FastAPI 部署训练好的模型为Web服务
"""

import os
import sys
import json
import yaml
import argparse
from pathlib import Path
from typing import Dict, List, Any, Optional
import logging
from datetime import datetime
import base64
from io import BytesIO

from fastapi import FastAPI, File, UploadFile, HTTPException, Form
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import uvicorn
from PIL import Image

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# API模型定义
class ExtractRequest(BaseModel):
    image_base64: str
    extract_fields: Optional[List[str]] = None


class ExtractResponse(BaseModel):
    success: bool
    extracted_fields: Dict[str, str]
    confidence_scores: Optional[Dict[str, float]] = None
    processing_time: float
    timestamp: str


class UniversalDocumentService:
    """通用文档信息抽取服务
    
    支持多种文档类型的信息抽取服务：
    - 合同文档
    - 发票文档
    - 简历文档
    - 医疗报告
    - 身份证
    - 其他结构化文档
    """
    
    def __init__(self, model_path: str, config_path: str = "config/config.yaml", document_type: str = "contract", gpu_id: Optional[int] = None):
        """初始化服务
        
        Args:
            model_path: 模型路径
            config_path: 配置文件路径
            document_type: 文档类型
            gpu_id: 指定GPU ID
        """
        # 设置GPU ID
        if gpu_id is not None:
            os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
            logger.info(f"🔧 已设置CUDA_VISIBLE_DEVICES={gpu_id}")
        
        self.model_path = self._resolve_model_path(model_path)
        self.adapter_path = None  # 初始化适配器路径
        self.config = self._load_config(config_path)
        self.document_type = document_type
        
        # 加载文档模板
        self.template = self._load_document_template(document_type)
        self.document_fields = [field['field_name'] for field in self.template['fields']]
        
        logger.info(f"📄 文档类型: {self.template['display_name']}")
        logger.info(f"🏷️  提取字段: {len(self.document_fields)} 个")
        logger.info(f"💾 模型路径: {self.model_path}")
        
        # 初始化模型
        self._setup_model()
    
    def _resolve_model_path(self, model_path: str) -> str:
        """解析模型路径，查找最新的checkpoint"""
        model_dir = Path(model_path)
        logger.info(f"🔍 解析模型路径: {model_path}")
        
        if model_dir.is_file():
            # 如果是文件，直接返回
            return str(model_dir)
            
        if not model_dir.exists():
            logger.warning(f"⚠️  模型路径不存在: {model_path}，将使用Swift CLI")
            return model_path
            
        # 查找最新的版本目录
        version_dirs = sorted([d for d in model_dir.iterdir() if d.is_dir() and d.name.startswith('v')])
        if version_dirs:
            latest_version = version_dirs[-1]
            logger.info(f"📁 找到最新版本: {latest_version.name}")
            
            # 检查是否有config.json文件（完整模型）
            config_file = latest_version / "config.json"
            if config_file.exists():
                logger.info(f"✅ 使用完整模型: {latest_version}")
                return str(latest_version)
            
            # 检查是否有checkpoint目录（完整模型）
            checkpoints = sorted(latest_version.glob('checkpoint-*'))
            if checkpoints:
                checkpoint_dir = checkpoints[-1]
                config_file = checkpoint_dir / "config.json"
                if config_file.exists():
                    logger.info(f"✅ 使用checkpoint: {checkpoint_dir}")
                    return str(checkpoint_dir)
                else:
                    # 检查是否为LoRA适配器模型
                    adapter_config = checkpoint_dir / "adapter_config.json"
                    if adapter_config.exists():
                        logger.info(f"🔧 检测到LoRA适配器模型: {checkpoint_dir}")
                        # 查找基础模型路径
                        with open(adapter_config, 'r', encoding='utf-8') as f:
                            adapter_cfg = json.load(f)
                        
                        base_model_path = adapter_cfg.get('base_model_name_or_path')
                        if base_model_path and Path(base_model_path).exists():
                            logger.info(f"📂 使用基础模型路径: {base_model_path}")
                            # 设置适配器路径
                            self.adapter_path = str(checkpoint_dir)
                            return base_model_path
                        else:
                            # 尝试在项目目录中查找基础模型
                            project_base_model = Path("models/universal_document/stage1") / latest_version.name / "checkpoint-3"
                            if project_base_model.exists() and (project_base_model / "config.json").exists():
                                logger.info(f"📂 使用项目基础模型: {project_base_model}")
                                self.adapter_path = str(checkpoint_dir)
                                return str(project_base_model)
                            else:
                                logger.warning(f"⚠️  无法找到基础模型，可能无法用于vLLM推理")
                                return str(checkpoint_dir)
        
        logger.info(f"📁 未找到版本目录，使用原路径: {model_path}")
        return model_path
        
    def _load_config(self, config_path: str) -> dict:
        """加载配置文件"""
        with open(config_path, 'r', encoding='utf-8') as f:
            return yaml.safe_load(f)
    
    def _load_document_template(self, document_type: str) -> dict:
        """加载文档模板配置"""
        template_path = Path(f"config/templates/{document_type}.yaml")
        
        if not template_path.exists():
            # 如果没有模板文件，使用默认合同模板
            logger.warning(f"⚠️  未找到模板文件: {template_path}，使用默认合同模板")
            return self._get_default_contract_template()
        
        try:
            with open(template_path, 'r', encoding='utf-8') as f:
                template = yaml.safe_load(f)
                logger.info(f"✅ 加载模板: {template['display_name']}")
                return template
        except Exception as e:
            logger.error(f"❌ 加载模板失败: {e}，使用默认合同模板")
            return self._get_default_contract_template()
    
    def _get_default_contract_template(self) -> dict:
        """获取默认合同模板"""
        return {
            'document_type': 'contract',
            'display_name': '合同文档',
            'description': '合同类文档的关键信息抽取',
            'fields': [
                {'field_name': '合同名称', 'field_type': 'text', 'required': True},
                {'field_name': '合同编号', 'field_type': 'text', 'required': True},
                {'field_name': '甲方名称', 'field_type': 'text', 'required': True},
                {'field_name': '甲方地址', 'field_type': 'text', 'required': False},
                {'field_name': '甲方联系人', 'field_type': 'text', 'required': False},
                {'field_name': '乙方名称', 'field_type': 'text', 'required': True},
                {'field_name': '乙方地址', 'field_type': 'text', 'required': False},
                {'field_name': '乙方联系人', 'field_type': 'text', 'required': False},
                {'field_name': '合同金额', 'field_type': 'amount', 'required': True},
                {'field_name': '履行期限', 'field_type': 'text', 'required': True},
                {'field_name': '签署日期', 'field_type': 'date', 'required': True},
                {'field_name': '生效日期', 'field_type': 'date', 'required': False},
                {'field_name': '履行地点', 'field_type': 'text', 'required': False},
                {'field_name': '付款方式', 'field_type': 'text', 'required': False},
                {'field_name': '违约责任', 'field_type': 'text', 'required': False}
            ]
        }
    
    def _setup_model(self):
        """设置模型推理服务"""
        try:
            # 首先尝试使用vLLM部署
            self._setup_vllm_service()
            
        except Exception as e:
            logger.warning(f"vLLM设置失败: {e}")
            # fallback到Swift推理
            self.use_vllm = False
            logger.info("将使用Swift CLI进行推理")
    
    def _setup_vllm_service(self):
        """设置vLLM推理服务"""
        try:
            from vllm import LLM, SamplingParams
            
            logger.info("🚀 初始化vLLM推理引擎...")
            
            # 检查配置文件
            config_file = Path(self.model_path) / "config.json"
            if not config_file.exists():
                raise FileNotFoundError(f"vLLM需要config.json文件，未在{self.model_path}中找到")
            
            # 准备vLLM初始化参数
            vllm_kwargs = {
                "model": self.model_path,
                "trust_remote_code": True,
                "gpu_memory_utilization": self.config['deployment']['vllm_gpu_memory_utilization'],
                "max_model_len": self.config['deployment']['vllm_max_model_len'],
                "enforce_eager": True,  # 避免图编译问题
                "disable_custom_all_reduce": True,  # 提高稳定性
                "dtype": "float16",  # 使用float16减少内存使用
                "tensor_parallel_size": 1,  # 单GPU
                # 允许本地媒体路径
                "allowed_local_media_path": "/home/ai/test_qwen/contract_extraction"
            }
            
            # 如果提供了适配器路径，则启用LoRA支持
            if self.adapter_path:
                logger.info(f"🔧 启用LoRA适配器支持: {self.adapter_path}")
                
                # 读取适配器配置获取基础模型路径
                adapter_config_path = os.path.join(self.adapter_path, "adapter_config.json")
                if os.path.exists(adapter_config_path):
                    with open(adapter_config_path, 'r', encoding='utf-8') as f:
                        adapter_config = json.load(f)
                    
                    base_model_path = adapter_config.get('base_model_name_or_path')
                    if base_model_path and os.path.exists(base_model_path):
                        logger.info(f"📂 使用基础模型路径: {base_model_path}")
                        vllm_kwargs["model"] = base_model_path
                    
                # 启用LoRA适配器
                vllm_kwargs.update({
                    "enable_lora": True,
                    "max_loras": 1,
                    "max_lora_rank": 64,
                    "lora_extra_vocab_size": 256
                })
            
            # 初始化vLLM模型
            self.vllm_model = LLM(**vllm_kwargs)
            
            # 设置采样参数（优化参数以提高准确性）
            self.sampling_params = SamplingParams(
                temperature=0.05,  # 降低温度以获得更确定的结果
                top_p=self.config['inference']['top_p'],
                max_tokens=self.config['inference']['max_new_tokens'],
                repetition_penalty=1.05,
                stop_token_ids=None
            )
            
            self.use_vllm = True
            logger.info("✅ vLLM模型加载成功")
            
        except ImportError:
            logger.info("📌 vLLM未安装，将使用Swift CLI")
            raise
        except Exception as e:
            logger.info(f"📌 vLLM设置失败: {e}")
            logger.info("🔄 转为使用Swift CLI推理（推荐）")
            raise
    
    def _create_extraction_prompt(self) -> str:
        """创建信息抽取提示词"""
        if 'prompt_template' in self.template:
            # 使用模板中的提示词
            field_list = "\n".join([f"- {field['field_name']}" for field in self.template['fields']])
            json_format = "{\n" + "\n".join([
                f'  "{field["field_name"]}": "具体内容或未找到"{"" if i == len(self.template["fields"]) - 1 else ","}'
                for i, field in enumerate(self.template['fields'])
            ]) + "\n}"
            
            return self.template['prompt_template'].format(
                field_list=field_list,
                json_format=json_format
            )
        else:
            # 使用优化的默认提示词
            prompt = f"""您是一位专业的{self.template['display_name']}审查专家，请仔细分析这份{self.template['display_name']}图像，准确提取以下关键字段信息。

提取要求：
1. 必须从图像中的{self.template['display_name']}准确提取信息
2. 如果某个字段在文档中找不到，请标记为"未找到"
3. 金额需要包含货币符号（如￥、$等）
4. 日期需要使用标准格式（如2024年1月15日或2024-01-15）
5. 公司名称需要完整准确

需要提取的字段：
"""
            for field in self.template['fields']:
                prompt += f"- {field['field_name']}\n"
            
            prompt += """
请严格按照以下JSON格式输出结果，不要添加任何其他内容：
{
"""
            for i, field in enumerate(self.template['fields']):
                comma = "," if i < len(self.template['fields']) - 1 else ""
                prompt += f'  "{field["field_name"]}": "具体内容或未找到"{comma}\n'
            
            prompt += "}"
            
            return prompt
    
    def _extract_with_vllm(self, image_base64: str, fields: Optional[List[str]] = None) -> Dict[str, str]:
        """使用vLLM进行信息抽取"""
        try:
            # 如果指定了字段，则创建临时模板
            if fields:
                # 创建临时模板
                temp_template = {
                    'document_type': self.template['document_type'],
                    'display_name': self.template['display_name'],
                    'description': self.template['description'],
                    'fields': [field for field in self.template['fields'] if field['field_name'] in fields]
                }
                
                # 使用临时模板创建提示词
                if 'prompt_template' in temp_template:
                    field_list = "\n".join([f"- {field['field_name']}" for field in temp_template['fields']])
                    json_format = "{\n" + "\n".join([
                        f'  "{field["field_name"]}": "具体内容或未找到"{"" if i == len(temp_template["fields"]) - 1 else ","}'
                        for i, field in enumerate(temp_template['fields'])
                    ]) + "\n}"
                    
                    prompt = temp_template['prompt_template'].format(
                        field_list=field_list,
                        json_format=json_format
                    )
                else:
                    # 使用默认提示词格式
                    prompt = f"""您是一位专业的{temp_template['display_name']}审查专家，请仔细分析这份{temp_template['display_name']}图像，准确提取以下关键字段信息。

提取要求：
1. 必须从图像中的{temp_template['display_name']}准确提取信息
2. 如果某个字段在文档中找不到，请标记为"未找到"
3. 金额需要包含货币符号（如￥、$等）
4. 日期需要使用标准格式（如2024年1月15日或2024-01-15）
5. 公司名称需要完整准确

需要提取的字段：
"""
                    for field in temp_template['fields']:
                        prompt += f"- {field['field_name']}\n"
                    
                    prompt += """
请严格按照以下JSON格式输出结果，不要添加任何其他内容：
{
"""
                    for i, field in enumerate(temp_template['fields']):
                        comma = "," if i < len(temp_template['fields']) - 1 else ""
                        prompt += f'  "{field["field_name"]}": "具体内容或未找到"{comma}\n'
                    
                    prompt += "}"
            else:
                # 使用完整模板创建提示词
                prompt = self._create_extraction_prompt()
            
            # vLLM多模态推理格式
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"}},
                        {"type": "text", "text": prompt}
                    ]
                }
            ]
            
            # 执行推理
            outputs = self.vllm_model.chat(
                messages=messages,
                sampling_params=self.sampling_params,
                use_tqdm=False
            )
            result_text = outputs[0].outputs[0].text
            
            # 解析结果
            return self._parse_extraction_result(result_text)
            
        except Exception as e:
            logger.error(f"vLLM推理失败: {e}")
            import traceback
            traceback.print_exc()
            raise
    
    def _extract_with_swift(self, image_base64: str, fields: Optional[List[str]] = None) -> Dict[str, str]:
        """使用Swift CLI进行信息抽取"""
        import subprocess
        import tempfile
        
        try:
            # 保存临时图像
            with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp_file:
                # 解码base64图像
                image_bytes = base64.b64decode(image_base64)
                tmp_file.write(image_bytes)
                tmp_path = tmp_file.name
            
            # 如果指定了字段，则创建临时模板提示词
            if fields:
                # 创建临时模板
                temp_template = {
                    'document_type': self.template['document_type'],
                    'display_name': self.template['display_name'],
                    'description': self.template['description'],
                    'fields': [field for field in self.template['fields'] if field['field_name'] in fields]
                }
                
                # 创建提示词
                prompt = f"""您是一位专业的{temp_template['display_name']}审查专家，请仔细分析这份{temp_template['display_name']}图像，准确提取以下关键字段信息。

提取要求：
1. 必须从图像中的{temp_template['display_name']}准确提取信息
2. 如果某个字段在文档中找不到，请标记为"未找到"
3. 金额需要包含货币符号（如￥、$等）
4. 日期需要使用标准格式（如2024年1月15日或2024-01-15）
5. 公司名称需要完整准确

需要提取的字段：
"""
                for field in temp_template['fields']:
                    prompt += f"- {field['field_name']}\n"
                
                prompt += """
请严格按照以下JSON格式输出结果，不要添加任何其他内容：
{
"""
                for i, field in enumerate(temp_template['fields']):
                    comma = "," if i < len(temp_template['fields']) - 1 else ""
                    prompt += f'  "{field["field_name"]}": "具体内容或未找到"{comma}\n'
                
                prompt += "}"
            else:
                # 使用完整模板创建提示词
                prompt = self._create_extraction_prompt()
            
            # 执行Swift推理
            cmd = [
                'swift', 'infer',
                '--model_type', self.config['model']['model_type'],
                '--model', self.model_path,
                '--images', tmp_path,
                '--query', prompt,
                '--max_new_tokens', str(self.config['inference']['max_new_tokens']),
                '--temperature', str(self.config['inference']['temperature']),
                '--top_p', str(self.config['inference']['top_p'])
            ]
            
            logger.info(f"🚀 执行Swift CLI命令: {' '.join(cmd)}")
            
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                check=True
            )
            
            # 清理临时文件
            os.unlink(tmp_path)
            
            # 解析结果
            return self._parse_extraction_result(result.stdout.strip())
            
        except subprocess.CalledProcessError as e:
            logger.error(f"Swift CLI推理失败: {e}")
            logger.error(f"stderr: {e.stderr}")
            logger.error(f"stdout: {e.stdout}")
            # 清理临时文件
            if 'tmp_path' in locals():
                os.unlink(tmp_path)
            raise
        except Exception as e:
            logger.error(f"Swift CLI推理失败: {e}")
            # 清理临时文件
            if 'tmp_path' in locals():
                os.unlink(tmp_path)
            raise
    
    def _parse_extraction_result(self, result_text: str) -> Dict[str, str]:
        """解析抽取结果"""
        try:
            # 尝试直接解析JSON
            # 查找JSON部分
            json_start = result_text.find('{')
            json_end = result_text.rfind('}') + 1
            
            if json_start >= 0 and json_end > json_start:
                json_text = result_text[json_start:json_end]
                return json.loads(json_text)
            else:
                logger.warning("未找到有效的JSON格式")
                raise json.JSONDecodeError("No valid JSON found", result_text, 0)
            
        except json.JSONDecodeError:
            # 如果JSON解析失败，尝试提取结构化信息
            logger.warning("JSON解析失败，尝试文本解析")
            return self._extract_from_text(result_text)
    
    def _extract_from_text(self, text: str) -> Dict[str, str]:
        """从文本中提取结构化信息"""
        result = {}
        import re
        
        for field in self.document_fields:
            # 尝试多种模式匹配
            patterns = [
                rf'"{field}"[：:\s]*"([^"]*)"',
                rf'{field}[：:\s]*([^\n，,。.]+)',
                rf'"{field}"[：:\s]*([^\n，,。.]+)'
            ]
            
            found = False
            for pattern in patterns:
                match = re.search(pattern, text)
                if match:
                    result[field] = match.group(1).strip()
                    found = True
                    break
            
            if not found:
                result[field] = "未找到"
        
        return result
    
    def extract_from_image(self, image_base64: str, 
                          fields: Optional[List[str]] = None) -> ExtractResponse:
        """从图像中抽取信息
        
        Args:
            image_base64: base64编码的图像
            fields: 要抽取的字段列表（可选）
            
        Returns:
            抽取结果
        """
        start_time = datetime.now()
        
        try:
            # 执行抽取
            if self.use_vllm:
                extracted_fields = self._extract_with_vllm(image_base64, fields)
            else:
                extracted_fields = self._extract_with_swift(image_base64, fields)
            
            # 计算处理时间
            processing_time = (datetime.now() - start_time).total_seconds()
            
            return ExtractResponse(
                success=True,
                extracted_fields=extracted_fields,
                processing_time=processing_time,
                timestamp=datetime.now().isoformat()
            )
            
        except Exception as e:
            logger.error(f"信息抽取失败: {e}")
            import traceback
            traceback.print_exc()
            processing_time = (datetime.now() - start_time).total_seconds()
            
            return ExtractResponse(
                success=False,
                extracted_fields={},
                processing_time=processing_time,
                timestamp=datetime.now().isoformat()
            )


from contextlib import asynccontextmanager

# 全局服务实例
service: Optional[UniversalDocumentService] = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    global service
    
    # 启动时初始化服务
    model_path = os.environ.get('MODEL_PATH', 'models/final_model')
    config_path = os.environ.get('CONFIG_PATH', 'config/config.yaml')
    document_type = os.environ.get('DOCUMENT_TYPE', 'contract')
    gpu_id = os.environ.get('GPU_ID')
    gpu_id = int(gpu_id) if gpu_id is not None else None
    
    try:
        service = UniversalDocumentService(model_path, config_path, document_type, gpu_id)
        logger.info("抽取服务初始化成功")
    except Exception as e:
        logger.error(f"服务初始化失败: {e}")
        sys.exit(1)
    
    yield
    
    # 关闭时清理资源
    service = None
    logger.info("服务已关闭")

# 创建FastAPI应用
app = FastAPI(
    title="通用文档信息抽取API",
    description="基于Qwen2.5-VL的多文档类型信息抽取服务",
    version="2.0.0",
    lifespan=lifespan
)


@app.post("/extract", response_model=ExtractResponse)
async def extract_contract_info(request: ExtractRequest):
    """抽取信息API"""
    if not service:
        raise HTTPException(status_code=503, detail="服务未初始化")
    
    try:
        result = service.extract_from_image(
            request.image_base64,
            request.extract_fields
        )
        return result
    except Exception as e:
        logger.error(f"API调用失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/extract_upload")
async def extract_from_upload(
    file: UploadFile = File(...),
    extract_fields: Optional[str] = Form(None)
):
    """通过文件上传抽取信息"""
    if not service:
        raise HTTPException(status_code=503, detail="服务未初始化")
    
    try:
        # 读取上传的图像
        image_bytes = await file.read()
        image = Image.open(BytesIO(image_bytes))
        
        # 转换为base64
        buffer = BytesIO()
        image.save(buffer, format='JPEG')
        image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
        
        # 解析字段列表
        fields_list = None
        if extract_fields:
            fields_list = [field.strip() for field in extract_fields.split(',')]
        
        # 执行抽取
        result = service.extract_from_image(image_base64, fields_list)
        return result
        
    except Exception as e:
        logger.error(f"文件上传API调用失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/health")
async def health_check():
    """健康检查API"""
    return {"status": "healthy", "timestamp": datetime.now().isoformat()}


@app.get("/info")
async def service_info():
    """服务信息API"""
    if not service:
        raise HTTPException(status_code=503, detail="服务未初始化")
    
    return {
        "model_path": service.model_path,
        "document_type": service.document_type,
        "supported_fields": service.document_fields,
        "template_name": service.template['display_name'],
        "use_vllm": service.use_vllm if hasattr(service, 'use_vllm') else False,
        "adapter_path": service.adapter_path if hasattr(service, 'adapter_path') else None
    }


@app.post("/test")
async def test_extraction():
    """测试接口 - 使用模拟数据测试抽取功能"""
    if not service:
        raise HTTPException(status_code=503, detail="服务未初始化")
    
    try:
        # 创建一个简单的测试图像base64
        from PIL import Image, ImageDraw
        import io
        
        # 创建测试图像
        img = Image.new('RGB', (400, 300), color='white')
        draw = ImageDraw.Draw(img)
        draw.text((50, 50), '测试合同文档', fill='black')
        draw.text((50, 100), '合同编号: TEST001', fill='black')
        draw.text((50, 150), '甲方: 测试公司A', fill='black')
        draw.text((50, 200), '乙方: 测试公司B', fill='black')
        
        # 转为base64
        buffer = io.BytesIO()
        img.save(buffer, format='JPEG')
        image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
        
        # 模拟抽取（使用模拟数据）
        mock_result = {
            field: f"测试_{field}_值" for field in service.document_fields[:5]
        }
        for field in service.document_fields[5:]:
            mock_result[field] = "未找到"
        
        return ExtractResponse(
            success=True,
            extracted_fields=mock_result,
            processing_time=0.1,
            timestamp=datetime.now().isoformat()
        )
        
    except Exception as e:
        logger.error(f"测试接口失败: {e}")
        raise HTTPException(status_code=500, detail=f"测试失败: {str(e)}")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="通用文档信息抽取服务部署")
    parser.add_argument('--model', required=True, help='模型路径')
    parser.add_argument('--config', default='config/config.yaml', help='配置文件路径')
    parser.add_argument('--document_type', default='contract',
                       choices=['contract', 'invoice', 'resume', 'medical_report', 'id_card'],
                       help='文档类型')
    parser.add_argument('--host', default='0.0.0.0', help='服务主机')
    parser.add_argument('--port', type=int, default=8000, help='服务端口')
    parser.add_argument('--workers', type=int, default=1, help='工作进程数')
    parser.add_argument('--gpu_id', type=int, default=None, help='指定GPU ID')
    
    args = parser.parse_args()
    
    # 设置环境变量
    os.environ['MODEL_PATH'] = args.model
    os.environ['CONFIG_PATH'] = args.config
    os.environ['DOCUMENT_TYPE'] = args.document_type
    if args.gpu_id is not None:
        os.environ['GPU_ID'] = str(args.gpu_id)
    
    # 启动服务
    logger.info(f"启动通用文档信息抽取服务...")
    logger.info(f"模型路径: {args.model}")
    logger.info(f"文档类型: {args.document_type}")
    if args.gpu_id is not None:
        logger.info(f"GPU ID: {args.gpu_id}")
    logger.info(f"服务地址: http://{args.host}:{args.port}")
    
    uvicorn.run(
        app,
        host=args.host,
        port=args.port,
        workers=args.workers,
        log_level="info"
    )


if __name__ == "__main__":
    main()